#include "th25d_40ub.h"
#include <string.h>
#include "t_spi_m.h"

/* 私有函数声明 */
static void th25d_cs_control(th25d_40ub_t *dev, uint8_t state);
static t_flash_erro_t th25d_write_enable(th25d_40ub_t *dev);
static t_flash_erro_t th25d_wait_for_ready(th25d_40ub_t *dev, uint32_t timeout);

/* 设备初始化与销毁 */
th25d_40ub_t *th25d_40ub_create(void *arg,
                               spi_flash_read_f read_func,
                               spi_flash_write_f write_func,
                               spi_cs_write_func hw_func,
                               const t_cs_gpio_t *hw_gpio)
{
    th25d_40ub_t *dev = (th25d_40ub_t *)T_MALLOC(sizeof(th25d_40ub_t));
    if (!dev) return T_NULL;

    dev->arg = arg;
    dev->hw_gpio = hw_gpio;
    dev->hw_func = hw_func;
    dev->read_func = read_func;
    dev->write_func = write_func;

    /* 初始化CS引脚 */
    if (dev->hw_gpio && dev->hw_func) {
        dev->hw_func(dev->hw_gpio->cs_port, dev->hw_gpio->cs_pin, 1);
    }

    return dev;
}

void th25d_40ub_destroy(th25d_40ub_t *dev)
{
    if (dev) {
        T_FREE(dev);
    }
}

/* CS引脚控制 */
static void th25d_cs_control(th25d_40ub_t *dev, uint8_t state)
{
    if (dev->hw_gpio && dev->hw_func) {
        dev->hw_func(dev->hw_gpio->cs_port, dev->hw_gpio->cs_pin, state);
    }
}

/* 写使能 */
static t_flash_erro_t th25d_write_enable(th25d_40ub_t *dev)
{
    uint8_t cmd = TH25D_CMD_WRITE_ENABLE;
    
    th25d_cs_control(dev, 0);
    dev->write_func(dev->arg, &cmd, 1, 100);
    th25d_cs_control(dev, 1);
    
    return TH25D_40UB_OK;
}

/* 等待设备就绪 */
static t_flash_erro_t th25d_wait_for_ready(th25d_40ub_t *dev, uint32_t timeout)
{
    uint8_t status;
    uint8_t cmd = TH25D_CMD_READ_STATUS_REG;
    
    do {
        th25d_cs_control(dev, 0);
        dev->write_func(dev->arg, &cmd, 1, 100);
        dev->read_func(dev->arg, &status, 1, 100);
        th25d_cs_control(dev, 1);
        
        if (!(status & TH25D_STATUS_WIP)) {
            return TH25D_40UB_OK;
        }
    } while (1);
    
    return TH25D_40UB_ERR_TIMEOUT;
}

/* 设备识别与状态 */
t_flash_erro_t th25d_40ub_read_id(th25d_40ub_t *dev, uint8_t *mid, uint8_t *did)
{
    uint8_t cmd = TH25D_CMD_READ_IDENTIFICATION;
    uint8_t buf[3];
    
    th25d_cs_control(dev, 0);
    dev->write_func(dev->arg, &cmd, 1, 100);
    dev->read_func(dev->arg, buf, 3, 100);
    th25d_cs_control(dev, 1);
    
    if (mid) *mid = buf[0];
    if (did) *did = buf[1];
    
    return TH25D_40UB_OK;
}

t_flash_erro_t th25d_40ub_read_unique_id(th25d_40ub_t *dev, uint8_t uid[16])
{
    uint8_t cmd = TH25D_CMD_READ_UNIQUE_ID;
    uint8_t dummy[4];
    
    th25d_cs_control(dev, 0);
    dev->write_func(dev->arg, &cmd, 1, 100);
    dev->read_func(dev->arg, dummy, 4, 100); // 跳过dummy字节
    dev->read_func(dev->arg, uid, 16, 100);
    th25d_cs_control(dev, 1);
    
    return TH25D_40UB_OK;
}

/* 基本读写操作 */
t_flash_erro_t th25d_40ub_read(th25d_40ub_t *dev, uint32_t addr, uint8_t *data, uint32_t size)
{
    uint8_t cmd[4] = {
        TH25D_CMD_READ_DATA,
        (uint8_t)(addr >> 16),
        (uint8_t)(addr >> 8),
        (uint8_t)addr
    };
    
    th25d_cs_control(dev, 0);
    dev->write_func(dev->arg, cmd, 4, 100);
    dev->read_func(dev->arg, data, size, 100);
    th25d_cs_control(dev, 1);
    
    return TH25D_40UB_OK;
}

t_flash_erro_t th25d_40ub_write(th25d_40ub_t *dev, uint32_t addr, const uint8_t *data, uint32_t size)
{
    t_flash_erro_t ret = TH25D_40UB_OK;
    uint32_t offset = 0;
    uint32_t chunk_size;
    uint8_t cmd[4];
    
    while (offset < size) {
        /* 计算当前页剩余空间 */
        chunk_size = TH25D_PAGE_SIZE - (addr % TH25D_PAGE_SIZE);
        if (chunk_size > (size - offset)) {
            chunk_size = size - offset;
        }
        
        /* 写使能 */
        if ((ret = th25d_write_enable(dev)) != TH25D_40UB_OK) {
            break;
        }
        
        /* 发送页编程命令 */
        cmd[0] = TH25D_CMD_PAGE_PROGRAM;
        cmd[1] = (uint8_t)(addr >> 16);
        cmd[2] = (uint8_t)(addr >> 8);
        cmd[3] = (uint8_t)addr;
        
        th25d_cs_control(dev, 0);
        dev->write_func(dev->arg, cmd, 4, 100);
        dev->write_func(dev->arg, (uint8_t*)data + offset, chunk_size, 100);
        th25d_cs_control(dev, 1);
        
        /* 等待编程完成 */
        if ((ret = th25d_wait_for_ready(dev, 100)) != TH25D_40UB_OK) {
            break;
        }
        
        offset += chunk_size;
        addr += chunk_size;
    }
    
    return ret;
}

/* 擦除操作 */
t_flash_erro_t th25d_40ub_erase_4k_sector(th25d_40ub_t *dev, uint32_t addr)
{
    uint8_t cmd[4] = {
        TH25D_CMD_SECTOR_ERASE_4K,
        (uint8_t)(addr >> 16),
        (uint8_t)(addr >> 8),
        (uint8_t)addr
    };
    
    /* 写使能 */
    t_flash_erro_t ret = th25d_write_enable(dev);
    if (ret != TH25D_40UB_OK) return ret;
    
    /* 发送擦除命令 */
    th25d_cs_control(dev, 0);
    dev->write_func(dev->arg, cmd, 4, 100);
    th25d_cs_control(dev, 1);
    
    /* 等待擦除完成 */
    return th25d_wait_for_ready(dev, 500);
}

t_flash_erro_t th25d_40ub_chip_erase(th25d_40ub_t *dev)
{
    t_flash_erro_t ret;
    uint32_t addr;
    
    /* 循环擦除所有64KB块 */
    for (addr = 0; addr < TH25D_40UB_CAPACITY; addr += TH25D_BLOCK_SIZE_64K) {
        uint8_t cmd[4] = {
            TH25D_CMD_BLOCK_ERASE_64K,
            (uint8_t)(addr >> 16),
            (uint8_t)(addr >> 8),
            (uint8_t)addr
        };
        
        /* 写使能 */
        if ((ret = th25d_write_enable(dev)) != TH25D_40UB_OK) {
            return ret;
        }
        
        /* 发送擦除命令 */
        th25d_cs_control(dev, 0);
        dev->write_func(dev->arg, cmd, 4, 100);
        th25d_cs_control(dev, 1);
        
        /* 等待擦除完成 */
        if ((ret = th25d_wait_for_ready(dev, 500)) != TH25D_40UB_OK) {
            return ret;
        }
    }
    
    return TH25D_40UB_OK;
}

/* 其他函数实现类似... */

/* 安全寄存器操作 */
t_flash_erro_t th25d_40ub_erase_security_reg(th25d_40ub_t *dev, uint8_t reg_num)
{
    if (reg_num < 1 || reg_num > 3) {
        return TH25D_40UB_ERR_INVAL;
    }
    
    uint32_t addr = (reg_num - 1) * TH25D_SECURITY_REG_SIZE;
    uint8_t cmd[4] = {
        TH25D_CMD_ERASE_SECURITY_REG,
        (uint8_t)(addr >> 16),
        (uint8_t)(addr >> 8),
        (uint8_t)addr
    };
    
    /* 写使能 */
    t_flash_erro_t ret = th25d_write_enable(dev);
    if (ret != TH25D_40UB_OK) return ret;
    
    /* 发送擦除命令 */
    th25d_cs_control(dev, 0);
    dev->write_func(dev->arg, cmd, 4, 100);
    th25d_cs_control(dev, 1);
    
    /* 等待操作完成 */
    return th25d_wait_for_ready(dev, 500);
}

/* 高级控制功能 */
t_flash_erro_t th25d_40ub_software_reset(th25d_40ub_t *dev)
{
    uint8_t cmd;
    
    /* 发送Enable Reset命令 */
    cmd = TH25D_CMD_ENABLE_RESET;
    th25d_cs_control(dev, 0);
    dev->write_func(dev->arg, &cmd, 1, 100);
    th25d_cs_control(dev, 1);
    
    /* 发送Reset命令 */
    cmd = TH25D_CMD_RESET;
    th25d_cs_control(dev, 0);
    dev->write_func(dev->arg, &cmd, 1, 100);
    th25d_cs_control(dev, 1);
    
    return th25d_wait_for_ready(dev, 100);
}